from dataset.xml_reader import BoxHandler


import matplotlib.pyplot as plt
from PIL import Image, ImageDraw
from xml.sax import make_parser
import pandas as pd
import numpy as np
import os


def draw_img_box(img_path, xml_path):
    img_path = r'F:\FlyAI\TBDetection_FlyAI\data\input\TBDetection\image\162.jpg'
    xml_path = r'F:\FlyAI\TBDetection_FlyAI\data\input\TBDetection\xml\582.xml'

    img = Image.open(img_path)
    parser = make_parser()
    handler = BoxHandler()
    parser.setContentHandler(handler)
    parser.parse(xml_path)
    draw = ImageDraw.Draw(img)
    for box in handler.boxes:
        draw.rectangle(box, outline='red', width=3)
    plt.imshow(img)
    plt.show()


def box_stat():
    train_csv = pd.read_csv(r'F:\FlyAI\TBDetection\data\input\validation.csv')
    base_xml_path = r'F:\FlyAI\TBDetection\data\input'

    xml_paths = [os.path.join(base_xml_path, p) for p in train_csv['xml_path'].values]
    n_file = len(xml_paths)
    parser = make_parser()
    handler = BoxHandler()
    parser.setContentHandler(handler)

    all_boxes = np.zeros((0, 4))

    for i, xml_path in enumerate(xml_paths):
        parser.parse(xml_path)
        box = np.array(handler.boxes)
        all_boxes = np.vstack([all_boxes, box])

        if box.shape[0] == 0:
            print('empty')

    total_area = (all_boxes[:, 2] - all_boxes[:, 0]) * (all_boxes[:, 3] - all_boxes[:, 1])

    print(f'[n_file: {n_file}] '
          f'[mean_box_num: {all_boxes.shape[0] / n_file}] '
          f'[mean_area: {total_area.mean()}] '
          f'[mean_position: {all_boxes.mean(0)}]')

    return all_boxes, total_area


def draw_all_box(boxes):
    img_path = r'F:\FlyAI\TBDetection_FlyAI\data\input\TBDetection\image\162.jpg'

    img = Image.open(img_path)
    draw = ImageDraw.Draw(img)
    for box in boxes:
        draw.rectangle(list(box), outline='red', width=3)
    plt.imshow(img)
    plt.show()


def draw_box_distribution(boxes):
    df = pd.DataFrame(boxes)
    for i in range(4):
        plt.subplot(2, 2, i + 1)
        freq = list(df[0].value_counts(bins=100, sort=False))
        plt.bar(range(len(freq)), freq)
    plt.show()


if __name__ == '__main__':
    draw_all_box(box_stat()[0])
    # box_stat()
